import json
import os


def labelme_to_yolo(json_file_path, output_dir, class_dict):
    """
    将 LabelMe JSON 文件转换为 YOLO 格式的 TXT 文件（支持rectangle和point）

    参数:
        json_file_path: LabelMe JSON 文件的路径
        output_dir: 输出TXT文件的目录
        class_dict: 类别名称到ID的映射字典 (例如 {"heart_beat": 0, "E": 1, "A": 2})
    """
    # 读取JSON文件
    with open(json_file_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    img_width = data['imageWidth']
    img_height = data['imageHeight']
    shapes = data['shapes']

    txt_filename = os.path.splitext(os.path.basename(json_file_path))[0] + '.txt'
    txt_filepath = os.path.join(output_dir, txt_filename)

    # 定义点的“微小”尺寸（归一化值）
    point_min_size = 0.002  # 你可以根据需要调整这个值

    with open(txt_filepath, 'w', encoding='utf-8') as f:
        for shape in shapes:
            label = shape['label']
            points = shape['points']
            shape_type = shape.get('shape_type', 'polygon')

            # 确保标签在类别字典中
            if label not in class_dict:
                print(f"警告: 跳过未在 class_dict 中定义的标签 '{label}'")
                continue
            class_id = class_dict[label]

            # 处理不同的形状类型
            if shape_type == 'rectangle':
                # 矩形有两个点：左上角和右下角
                x_min, y_min = points[0]
                x_max, y_max = points[1]

            elif shape_type == 'point':
                # 点只有一个坐标 [x, y]
                # 将其视为一个极小的矩形框
                x_center = points[0][0]
                y_center = points[0][1]
                x_min = x_center - (img_width * point_min_size / 2)
                x_max = x_center + (img_width * point_min_size / 2)
                y_min = y_center - (img_height * point_min_size / 2)
                y_max = y_center + (img_height * point_min_size / 2)

            elif shape_type == 'polygon':
                # 提取所有x和y坐标
                x_coords = [p[0] for p in points]
                y_coords = [p[1] for p in points]
                x_min, x_max = min(x_coords), max(x_coords)
                y_min, y_max = min(y_coords), max(y_coords)
            else:
                # 跳过其他不支持的类型
                print(f"警告: 跳过不支持的形状类型 '{shape_type}'")
                continue

            # 计算YOLO格式的归一化中心坐标和宽高
            box_width = x_max - x_min
            box_height = y_max - y_min
            center_x = (x_min + x_max) / 2.0
            center_y = (y_min + y_max) / 2.0

            # 归一化
            center_x_norm = center_x / img_width
            center_y_norm = center_y / img_height
            box_width_norm = box_width / img_width
            box_height_norm = box_height / img_height

            # 确保归一化后的值在[0,1]范围内（防止由于点的小矩形超出图像边界）
            center_x_norm = max(0, min(1, center_x_norm))
            center_y_norm = max(0, min(1, center_y_norm))
            box_width_norm = max(0, min(1, box_width_norm))
            box_height_norm = max(0, min(1, box_height_norm))

            # 写入TXT文件
            f.write(f"{class_id} {center_x_norm:.6f} {center_y_norm:.6f} {box_width_norm:.6f} {box_height_norm:.6f}\n")


# 示例用法
# 注意：现在 class_dict 需要包含所有类型的标签（heart_beat, E, A）
class_dict = {"heart_beat": 0, "E": 1, "A": 2}
json_file_path = "D:\\project202509\\20250514_SpecParam_EA\\202407051519.json"  # 你的LabelMe JSON文件路径
output_dir = "D:\\Data\\20250514_SpecParam_EA\\labelme_to_yolo"  # 输出目录

# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)

# 转换单个文件
# labelme_to_yolo(json_file_path, output_dir, class_dict)
print(f"转换完成！TXT文件已保存至: {output_dir}")

# 如果要转换一个目录下的所有JSON文件
json_dir = "D:\Data\\20250514_SpecParam_EA"
for json_filename in os.listdir(json_dir):
    if json_filename.endswith('.json'):
        json_path = os.path.join(json_dir, json_filename)
        labelme_to_yolo(json_path, output_dir, class_dict)
        print(f"转换完成！TXT文件已保存至: {output_dir}")
